# _*_ coding: utf-8 _*_
"""
Time:     2024/9/14
Author:   2051261
File:     base_src_entity.py
"""
from __future__ import annotations
import typing as t
import logging

from schematics.models import Model
from schematics.types import StringType
from schematics.types import ListType
from schematics.types import BooleanType
from sqlalchemy.exc import OperationalError
import pymysql

from mine.db.db_operate import db_operate
from extractor.utils.common_func import dict_keys_to_lowercase
from extractor.const import ContextKey
from extractor.model.factory import SQLBuilderFactory
from extractor.flow.base_node_flow import BaseNodeFlow
from extractor.flow.node_model.child_node_method.modify import strip_space

logger = logging.getLogger(__name__)


if t.TYPE_CHECKING:
    from extractor.flow.node_model.log_node import RecordLog
    from extractor.flow.task_entry import Cache


class BaseSrcEntityModel(Model):
    table_name: str = StringType(required=True)
    conn_name: str = StringType(required=True)
    order_by: str = StringType(required=True)
    inc_column: str = StringType()
    operator: str = StringType(default=">")
    query_columns: list[str] = ListType(StringType())
    strip_space: bool = BooleanType(default=True)


class BaseSrcEntityNode(BaseNodeFlow):
    """
    kwargs 中的参数就是 BaseSrcEntityModel 中定义的
    """
    async def __call__(self, record_log: RecordLog, context, cache: Cache, record: t.DefaultDict, kwargs=None) \
            -> t.Iterator[t.Dict[str, t.Any]]:
        model = BaseSrcEntityModel(kwargs)
        cache_value = await cache.get_value()
        operator = model.operator

        query_columns = model.query_columns if model.query_columns else "*"
        conn_name, table_name = model.conn_name, model.table_name
        db_type = await db_operate.get_db_type(conn_name)
        sql_builder = await SQLBuilderFactory().get_query_builder(db_type)
        inc_column = model.inc_column
        order_by = model.order_by

        sql = sql_builder.query_sql(table_name, order_by, inc_column, query_columns, operator)
        result = db_operate.fetch_many(conn_name, sql, {inc_column: cache_value})
        try:
            for item in result:
                tmp = context.copy()
                # item = encode_item(item)
                if model.strip_space is True:
                    strip_space(item)
                item = await dict_keys_to_lowercase(item)
                tmp[ContextKey.DATA] = item
                yield tmp
        except OperationalError as e:
            if isinstance(e.orig, pymysql.err.OperationalError) and (e.orig.args[0] == 2013 or e.orig.args[0] == 2006):
                # 当数据源比较大清洗时间可能很长,丢失连接时,等待下次定时触发
                logger.warning(f"connection lost, table:{table_name}, e:{e}")
                return
            raise e
        except Exception as e:
            logger.error(f"get data from conn: {conn_name}, table:{conn_name}.{table_name} error, e:{e}")
            raise e


if __name__ == "__main__":
    data = {
        "table_name": "TB_COMPANY_0088",
        "conn_name": "cn_info_db",
        "order_by": "OB_SEQ_ID",
        "inc_column": "OB_SEQ_ID",
        "operator": ">",
        "query_columns": "*",
        "strip_space": True
    }

    BaseSrcEntityModel(data)
